from gelslimpi import GelSlimInterface, InputSource
from gelslimpi.visualize import qtcv, QImageViewWidget, QTimer, Qt
from gelslimpi.marker_tracker import *
from gelslimpi.utils import VideoReader, VirableRateVideoWriter, increment_path
from pyati import ATISensor
import cv2
import sys
from PyQt5 import QtGui, QtWidgets, QtCore
from scipy.interpolate import LinearNDInterpolator
from scipy.interpolate import RegularGridInterpolator as RGI

from datetime import datetime

from time import sleep
import numpy as np
from pathlib import Path
BASE_DIR = Path(__file__).resolve().parent
CONF_DIR = BASE_DIR/"videos"


def draw_marker(img, dots):
    img = img.copy()
    row, col = dots.shape[:2]
    dots = dots.reshape(-1, 2)
    for i, dot in enumerate(dots):
        #color 为hsv渐变色
        # color = np.array([[[i / col * 180/7, 255, 255]]], dtype=np.uint8)
        # color = cv2.cvtColor(color, cv2.COLOR_HSV2RGB)[0][0]
        cv2.circle(img, (int(dot[0]), int(dot[1])), 10, (0, 225, 225), 1)
    return img


class MarkerCollector(QtWidgets.QWidget):
    """移动到任意像素坐标
    采集图片
    采集传感器数据
    下压"""
    def __init__(self, 
        log_dir = BASE_DIR / f"runs/{datetime.now().strftime('%y_%m_%d_%H_%M_%S')}",
        input = BASE_DIR/"videos/video_0003.mp4",  # "/raspicam_node32/image/compressed"
        input_type = "video",        
        configs = CONF_DIR/"configs.yaml",
        rectify_on = False,
        save = False,
    ):
        super().__init__()
        self.save = save
        self.ati = ATISensor()
        self.ati.tare()
        self.source = InputSource(input, input_type)
        if input_type == "video":
            self.n_frame = self.source.input_source.total_frames
            self.duration = self.source.input_source.duration
        self.cnt = 1

        self.gelslim = GelSlimInterface(configs=configs, input_source=self.source)
        self.gelslim.configs["init"]["rectify_on"] = rectify_on
        self.init_ui()
        self.timer = QtCore.QTimer(self)
        self.timer.start(50)
        self.timer.timeout.connect(self.on_timeout)
        self.image_view.sigMousePressed.connect(self.on_mouse)
        
        # marker tracker
        img = self.get_frame()
        dots = self.gelslim.marker_detect(self.gelslim.marker_binarize(img))
        self.marker_tracker = MarkerTracker(dots, self.gelslim.configs["marker_track"])
        self.start_pts = self.marker_tracker.start_pts

        # save data
        self.dots = []
        self.ids = []
        self.forces = []
        self.log_dir = Path(log_dir)
        if save:
            self.log_dir.mkdir(exist_ok=True, parents=True)
        self.last_dots = self.marker_tracker.start_pts
        self.is_recording = False
        self.mouse_flag = False

    def init_ui(self):
        self.resize(1280, 960)
        hbox = QtWidgets.QHBoxLayout(self)
        self.image_view = QImageViewWidget(self)
        hbox.addWidget(self.image_view)

    def on_timeout(self):
        img = self.get_frame()
        force = self.get_ati()
        img_bin = self.gelslim.marker_binarize(img)
        
        self.marker_tracker.update_dots(img_bin)
        self.marker_tracker.update_dots(img_bin)
        self.marker_tracker.update_dots(img_bin)
        self.marker_tracker.update_dots(img_bin)

        img_show = draw_marker(img, self.marker_tracker.curr_pts)

        if self.mouse_flag:
            self.mouse_flag = False
            self.record_data(img, self.marker_tracker.curr_pts, force)
        if self.cnt % int(self.n_frame / int(self.duration*5)) == 0 and self.save:
            self.record_data(img, self.marker_tracker.curr_pts, force)
        if self.cnt >= self.n_frame:
            self.close()
        self.cnt += 1
        # if self.is_recording:
        #     self.auto_record_data(img, self.marker_tracker.curr_pts, force, img_show)

        self.image_view.setData(img_show)

    def keyPressEvent(self, ev) -> None:
        """按键控制
        """
        key = ev.key()
        if key in [Qt.Key.Key_Return, Qt.Key.Key_Enter]:
            self.is_recording = not self.is_recording
            if not self.is_recording:
                np.savez(str(self.log_dir/"data.npz"))
    
    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
        if len(self.dots)>0:
            np.savez(str(self.log_dir/"data.npz"), 
                dots=np.array(self.dots), 
                id=np.array(self.ids))
        return super().closeEvent(a0)

    def get_frame(self):
        return self.gelslim.get_frame()[0]

    def get_ati(self):
        return self.ati.data

    def auto_record_data(self, img, dots, force, img_show):
        tmp = np.sum(np.sum(np.abs(dots - self.last_dots), axis=2) > 8)
        if tmp < 3:
            return 
        self.last_dots = dots.copy()
        self.dots.append(dots.copy())
        self.forces.append(force)
        img_path = str(self.log_dir / f"{len(self.forces)-1:04d}.jpg")
        # imgshow_path = str(self.log_dir / f"vis_{len(self.forces)-1:04d}.jpg")
        cv2.imwrite(img_path, img)
        # cv2.imwrite(imgshow_path, img_show)
        print(f"{tmp}, save {img_path}")
    

    def record_data(self, img, dots, force):
        diff = np.abs(self.start_pts - dots)
        print(np.mean(diff))
        if np.random.uniform(0, 5) > np.mean(diff) or np.random.uniform(0, 15)>np.max(diff)  :
            return 
        self.last_dots = dots.copy()
        self.dots.append(dots.copy())
        self.forces.append(force)
        id = int(len(self.dots)-1)
        self.ids.append(id)
        img_path = str(self.log_dir / f"{id:04d}.jpg")
        cv2.imwrite(img_path, img)
        # imgshow_path = str(self.log_dir / f"vis_{len(self.forces)-1:04d}.jpg")
        # cv2.imwrite(imgshow_path, img_show)
        print(f"save {img_path}")

    def on_mouse(self, xy):
        self.mouse_flag = True


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = MarkerCollector(
            input = BASE_DIR/"videos/zonghe.mp4", 
            rectify_on=False,
            save=False
        )
    win.show()
    sys.exit(app.exec_())

